# Generate dataset.json for nnU-Net
import os, json

dataset_dir = "/data/suziren25/suziren/nnUNet/nnUNet_raw/Dataset506_KidneyUS"
imagesTr = os.path.join(dataset_dir, "imagesTr")
labelsTr = os.path.join(dataset_dir, "labelsTr")

cases = sorted([f.replace("_0000.nii.gz", "") for f in os.listdir(imagesTr) if f.endswith("_0000.nii.gz")])

dataset_json = {
    "name": "KidneyUS",
    "description": "Kidney segmentation in abdominal ultrasound images",
    "tensorImageSize": "2D",
    "modality": {
        "0": "US"
    },
    "channel_names": {
        "0": "ultrasound"
    },
    "labels": {
        "background": 0,
        "kidney": 1
    },
    "file_ending": ".nii.gz",
    "numTraining": len(cases),
    "numTest": 0,
    "training": [
        {
            "image": f"./imagesTr/{c}_0000.nii.gz",
            "label": f"./labelsTr/{c}.nii.gz"
        }
        for c in cases
    ],
    "test": []
}

with open(os.path.join(dataset_dir, "dataset.json"), "w") as f:
    json.dump(dataset_json, f, indent=4)

print(f"✅ dataset.json 已更新，共包含 {len(cases)} 个训练样本")
